#include <ntifs.h>
#include <ntstrsafe.h>
#include <windef.h>

#include "CRegistry.h"
#include "AllocateMem.h"


#ifdef DBG
#define Log(...) KdPrint((__VA_ARGS__))
#else
#define Log(...)
#endif

namespace wk
{
CRegistryObject::CRegistryObject()
    : m_ustrKeyPath({0})
    , m_hKey(NULL)
{
    Log("[Debug] CRegistryObject::CRegistryObject()\r\n");
}

CRegistryObject::CRegistryObject(
    _In_ LPCWSTR pwsKeyPath
    )
{
    Log("[Debug] CRegistryObject::CRegistryObject(LPCWSTR)\r\n");
    OpenKey(pwsKeyPath);
}

CRegistryObject::CRegistryObject(
    _In_ PUNICODE_STRING pustrKeyPath
    )
{
    Log("[Debug] CRegistryObject::CRegistryObject(PUNICODE_STRING)\r\n");
    OpenKey(pustrKeyPath);
}

CRegistryObject::~CRegistryObject()
{
    Log("[Debug] CRegistryObject::~CRegistryObject()\r\n");

    Destroy();
}

VOID CRegistryObject::Init(
    _In_ LPWSTR pwsKeyPath, 
    _In_ HANDLE hKey
    )
{
    Log("[Debug] CRegistryObject::Init(LPWSTR, HANDLE)\r\n");

    if (pwsKeyPath && hKey)
    {
        RtlCreateUnicodeString(&m_ustrKeyPath, pwsKeyPath);
        m_hKey = hKey;
    }
}

VOID CRegistryObject::Destroy()
{
    Log("[Debug] CRegistryObject::Destroy()\r\n");

    if (m_hKey)
    {
        Log("[Debug] CRegistryObject::Destroy m_hKey=%#X\r\n", m_hKey);
        ZwClose(m_hKey);
        m_hKey = NULL;
    }
    if ( m_ustrKeyPath.Buffer)
    {
        Log("[Debug] CRegistryObject::Destroy m_pustrKeyPath=%wZ\r\n", m_ustrKeyPath);
        RtlFreeUnicodeString(&m_ustrKeyPath);
        m_ustrKeyPath = { 0 };
    }
}

/*******************************************************************************
* Registry Key
*******************************************************************************/

NTSTATUS CRegistryObject::CreateKey(
    _In_ LPCWSTR pwsKeyPath
    )
{
    Log("[Debug] CRegistryObject::CreateKey(LPCWSTR)\r\n");

    NTSTATUS status;
    UNICODE_STRING ustrKeyName;
    OBJECT_ATTRIBUTES attributes;
    ULONG dwCreateResult;
    
    RtlInitUnicodeString(&ustrKeyName, pwsKeyPath);
    InitializeObjectAttributes(&attributes, &ustrKeyName, 0, NULL, NULL);
    status = ZwCreateKey(&m_hKey, KEY_ALL_ACCESS, &attributes, 0, NULL, REG_OPTION_NON_VOLATILE, &dwCreateResult);
    if (NT_SUCCESS(status))
    {
        RtlCreateUnicodeString(&m_ustrKeyPath, pwsKeyPath);
    }

    return status;
}

NTSTATUS CRegistryObject::OpenKey(
    _In_ LPCWSTR pwsKeyPath
    )
{
    Log("[Debug] CRegistryObject::OpenKey(LPCWSTR)\r\n");

    UNICODE_STRING ustrKeyPath;

    RtlInitUnicodeString(&ustrKeyPath, pwsKeyPath);
    return OpenKey(&ustrKeyPath);
}

NTSTATUS CRegistryObject::OpenKey(
    _In_ PUNICODE_STRING pustrKeyPath
    )
{
    Log("[Debug] CRegistryObject::OpenKey(PUNICODE_STRING)\r\n");

    NTSTATUS status;
    OBJECT_ATTRIBUTES attributes;

    InitializeObjectAttributes(&attributes, pustrKeyPath, 0, NULL, NULL);
    status = ZwOpenKey(&m_hKey, KEY_ALL_ACCESS, &attributes);
    if (NT_SUCCESS(status))
    {
        RtlCreateUnicodeString(&m_ustrKeyPath, pustrKeyPath->Buffer);
    }

    return status;
}

NTSTATUS CRegistryObject::DeleteKey()
{
    Log("[Debug] CRegistryObject::DeleteKey()\r\n");

    NTSTATUS status;

    if (m_hKey)
    {
        status = ZwDeleteKey(m_hKey);
        Destroy();
        return status;
    }
    return STATUS_UNSUCCESSFUL;
}


/*******************************************************************************
* Registry SubKey
*******************************************************************************/

NTSTATUS CRegistryObject::CreateSubKey(
    _In_ LPWSTR pwsSubKeyName,
    _Out_ CRegistryObject* pSubKeyRegObj
    )
{
    KdPrint(("[Debug] CRegistryObject::CreateSubKey(LPWSTR, CRegistryObject*)\r\n"));

    NTSTATUS status;
    HANDLE hKey;
    UNICODE_STRING ustrSubKeyName;
    OBJECT_ATTRIBUTES attributes;
    ULONG dwCreateResult;
    
    RtlInitUnicodeString(&ustrSubKeyName, pwsSubKeyName);
    InitializeObjectAttributes(&attributes, &ustrSubKeyName, 0, m_hKey, NULL);
    status = ZwCreateKey(&hKey, KEY_ALL_ACCESS, &attributes, 0, NULL, REG_OPTION_NON_VOLATILE, &dwCreateResult);
    if (pSubKeyRegObj && NT_SUCCESS(status))
    {
        WCHAR wszSubKeyPath[MAX_PATH];
        RtlStringCbPrintfW(wszSubKeyPath, sizeof(wszSubKeyPath), L"%s\\%s", m_ustrKeyPath.Buffer, pwsSubKeyName);
        KdPrint(("[Debug] CRegistryObject::CreateSubKey szSubKeyPath=%ws\r\n", wszSubKeyPath));
        pSubKeyRegObj->Init(wszSubKeyPath, hKey);
    }
    return status;
}

NTSTATUS CRegistryObject::OpenSubKey(
    _In_ LPWSTR pwsSubKeyName,
    _Out_ CRegistryObject* pSubKeyRegObj
    )
{
    KdPrint(("[Debug] CRegistryObject::OpenSubKey(LPWSTR, CRegistryObject*)\r\n"));

    UNICODE_STRING ustrSubKeyName;
    RtlInitUnicodeString(&ustrSubKeyName, pwsSubKeyName);
    return OpenSubKey(&ustrSubKeyName, pSubKeyRegObj);
}

NTSTATUS CRegistryObject::OpenSubKey(
    _In_ PUNICODE_STRING pustrSubKeyName,
    _Out_ CRegistryObject* pSubKeyRegObj
    )
{
    KdPrint(("[Debug] CRegistryObject::OpenSubKey(PUNICODE_STRING, CRegistryObject*)\r\n"));

    WCHAR wszRegPath[MAX_PATH];
    UNICODE_STRING ustrRegPath;
    
    if (pSubKeyRegObj)
    {
        RtlInitEmptyUnicodeString(&ustrRegPath, wszRegPath, sizeof(wszRegPath));
        RtlUnicodeStringPrintf(&ustrRegPath, L"%s\\%s", m_ustrKeyPath.Buffer, pustrSubKeyName->Buffer);
        return pSubKeyRegObj->OpenKey(&ustrRegPath);
    }
    
    return STATUS_INVALID_PARAMETER_2;
}

NTSTATUS CRegistryObject::DeleteSubKey(
    _In_ LPWSTR pwsSubKeyName
    )
{
    KdPrint(("[Debug] CRegistryObject::DeleteSubKey(LPWSTR)\r\n"));

    NTSTATUS status;
    CRegistryObject regObj;

    status = OpenSubKey(pwsSubKeyName, &regObj);
    if (!NT_SUCCESS(status))
    {
        return status;
    }
    return regObj.DeleteKey();
}


/*******************************************************************************
* Registry Value
*******************************************************************************/


}

